//
// Created by steven on 2025/9/4.
//

#ifndef MAIN_DETECTOR_CALIBRATOR_H
#define MAIN_DETECTOR_CALIBRATOR_H

typedef struct DataFlowParams
{
  std::size_t batchSize;
  std::size_t imageIndex;
  std::size_t maxBatches;
  std::string imgListFilePath;
  bool readCache;
  std::string calibratorTablePath;
} DataFlowParams;

std::vector<float> preprocessImage(const cv::Mat& image)
{
    cv::Mat processed;
    int channels =3;
    int height = 640;
    int width = 640;

    // 调整大小
    cv::resize(image, processed, cv::Size(width, height));

    // 转换颜色空间
    if (channels == 3) {
        cv::cvtColor(processed, processed, cv::COLOR_BGR2RGB);
    }
    else if (channels == 1) {
        cv::cvtColor(processed, processed, cv::COLOR_BGR2GRAY);
    }

    // 转换为浮点数并归一化
    processed.convertTo(processed, CV_32F, 1.0 / 255.0);

    // 转换为CHW格式并调整为模型输入范围
    std::vector<float> data;
    data.reserve(channels * height * width);

    if (channels == 3) {
        for (int c = 0; c < 3; ++c) {
            for (int h = 0; h < height; ++h) {
                for (int w = 0; w < width; ++w) {
                    data.push_back(processed.at<cv::Vec3f>(h, w)[c]);
                }
            }
        }
    }
    else if (channels == 1) {
        for (int h = 0; h < height; ++h) {
            for (int w = 0; w < width; ++w) {
                data.push_back(processed.at<float>(h, w));
            }
        }
    }

    return data;
}

class Int8EntropyCalibrator : public TopsInference::IInt8EntropyCalibrator
{
public:
  explicit Int8EntropyCalibrator(const DataFlowParams& data_params);
  virtual ~Int8EntropyCalibrator();
  int getBatchSize() const noexcept override { return data_params_.batchSize; }
  bool getBatch(TopsInference::TensorPtr_t bindings[], const char* names[],
                int num) noexcept override;
  const void* readCalibrationCache(int64_t& length) override;
  bool writeCalibrationCache(const void* ptr, int64_t length) override;

private:
  std::size_t batch_index_;
  DataFlowParams data_params_;
  std::size_t input_size_;
  std::vector<std::string> img_paths_;
  std::vector<float> h_inputs_;
  std::vector<char> calibrationCache;
};

Int8EntropyCalibrator::Int8EntropyCalibrator(const DataFlowParams& data_params)
  : data_params_(data_params)
{
  batch_index_ = data_params_.imageIndex +
    (data_params_.batchSize - 1) / data_params_.batchSize;
  int inputChannel = 3;
  int inputH = 640;
  int inputW = 640;
  this->input_size_ = data_params_.batchSize * inputChannel * inputH * inputW;
  std::fstream f(data_params.imgListFilePath, std::ios_base::in);
  if (f.is_open()) {
    std::string temp;
    while (std::getline(f, temp)) img_paths_.push_back(temp);
  }
  h_inputs_.resize(this->input_size_);
}

Int8EntropyCalibrator::~Int8EntropyCalibrator() {}

bool Int8EntropyCalibrator::getBatch(TopsInference::TensorPtr_t bindings[],
                                     const char* names[],
                                     int nbBindings) noexcept
{
  if (batch_index_ > data_params_.maxBatches ||
    data_params_.imageIndex + data_params_.batchSize >
    static_cast<int>(img_paths_.size())) {
    return false;
  }
  float* h_inputs_block = h_inputs_.data();
  std::size_t img_ind_st = data_params_.imageIndex;
  std::size_t img_ind_ed = data_params_.imageIndex + data_params_.batchSize;
  for (size_t j = img_ind_st; j < img_ind_ed; ++j) {
    if (j % 5 == 0 || j == img_ind_ed) {
      std::cout << "load img percent: " << img_paths_[j]
        << " +++++++++++++++: " << (j + 1) * 100. / img_paths_.size()
        << "%" << std::endl;
    }
    cv::Mat img = cv::imread(img_paths_[j]);
    std::vector<float> inputData = preprocessImage(img);
    if (inputData.size() != input_size_) {
      std::cerr << "input shape error" << std::endl;
      return false;
    }
    assert(inputData.size() == input_size_);
    std::memcpy(h_inputs_block, inputData.data(),
                static_cast<int>(inputData.size()) * sizeof(float));
    h_inputs_block += inputData.size();
  }
  data_params_.imageIndex += data_params_.batchSize;
  assert(nbBindings == 1);
  TopsInference::TensorPtr_t sub_input = bindings[0];
  sub_input->setOpaque(reinterpret_cast<void*>(h_inputs_.data()));
  sub_input->setDeviceType(TopsInference::DataDeviceType::HOST);
  TopsInference::Dims input_shape;
  input_shape.nbDims = 4;
  input_shape.dimension[0] = data_params_.batchSize; // N
  input_shape.dimension[1] = 3; // C
  input_shape.dimension[2] = 640; // H
  input_shape.dimension[3] = 640; // W
  sub_input->setDims(input_shape);
  batch_index_++;
  return true;
}

const void* Int8EntropyCalibrator::readCalibrationCache(int64_t& length)
{
  std::cout << __FUNCTION__ << ":" << __LINE__ << std::endl;
  calibrationCache.clear();

  std::ifstream cacheFile(data_params_.calibratorTablePath, std::ios::binary);

  if (cacheFile.good())
  {
      cacheFile.seekg(0, std::ios::end);
      size_t size = cacheFile.tellg();
      cacheFile.seekg(0, std::ios::beg);

      calibrationCache.resize(size);
      cacheFile.read(calibrationCache.data(), size);
      length = size;
      return calibrationCache.data();
  }

  length = 0;
  return nullptr;
}

bool Int8EntropyCalibrator::writeCalibrationCache(const void* cache,
                                                  int64_t length)
{
  std::ofstream output(data_params_.calibratorTablePath, std::ios::binary);
  output.write(reinterpret_cast<const char*>(cache), length);
  return true;
}

#endif //MAIN_DETECTOR_CALIBRATOR_H